"""Create Package variants for PyPI distribution."""

import argparse
import os
import tomllib

from packaging.version import Version
from test_utils import PY_PACKAGE

IN_PATH = os.path.join(PY_PACKAGE, "pyproject.toml.in")
STUB_IN_PATH = os.path.join(PY_PACKAGE, "pyproject.toml.stub.in")
OUT_PATH = os.path.join(PY_PACKAGE, "pyproject.toml")

NCCL_WHL = """    \"nvidia-nccl-{0} ; platform_system == 'Linux' and platform_machine != 'aarch64'\","""

NAME = "{{ name }}"
NCCL = "{{ nccl }}"
VERSION = "{{ version }}"
CUDA_VARIANTS = ["cu12", "cu13"]


def copyfile(src: str, dst: str) -> None:
    with open(src, "rb") as fd:
        content = fd.read()
    with open(dst, "wb") as fd:
        fd.write(content)


def make_pyproject(
    *, use_suffix: str, require_nccl_dep: str, create_stub: bool = False
) -> None:
    if use_suffix == "cpu" and require_nccl_dep != "na":
        raise ValueError(
            "xgboost-cpu cannot require NCCL dependency. "
            "When setting --use-suffix='cpu', you must also set --require-nccl-dep='na'."
        )
    if (
        use_suffix in CUDA_VARIANTS
        and require_nccl_dep in CUDA_VARIANTS
        and use_suffix != require_nccl_dep
    ):
        raise ValueError(
            "Inconsistent choices for --use-suffix and --require-nccl-dep. "
            "When --use-suffix is set to one of {{{0}}}, --require-nccl-dep must be "
            "set to identical value as --use-suffix.".format(",".join(CUDA_VARIANTS))
        )
    if create_stub:
        if use_suffix == "na":
            raise ValueError("To create a stub package, --use-suffix must not be 'na'")
        if require_nccl_dep != "na":
            raise ValueError(
                "To create a stub package, --require-nccl-dep must be 'na'"
            )

    with open(STUB_IN_PATH if create_stub else IN_PATH) as fd:
        pyproject = fd.read()

    readme_dft = os.path.join(PY_PACKAGE, "README.dft.rst")
    readme_cpu = os.path.join(PY_PACKAGE, "README.cpu.rst")
    readme_stub = os.path.join(PY_PACKAGE, "README.stub.rst")
    readme = os.path.join(PY_PACKAGE, "README.rst")
    pyproject = pyproject.replace(
        NAME, f"xgboost-{use_suffix}" if use_suffix != "na" else "xgboost"
    )
    if create_stub:
        copyfile(readme_stub, readme)
        pyproject_parsed = tomllib.loads(pyproject)
        pyproject = pyproject.replace(
            VERSION, str(Version(pyproject_parsed["project"]["version"]))
        )
    elif use_suffix == "cpu":
        copyfile(readme_cpu, readme)
    else:
        copyfile(readme_dft, readme)
    pyproject = pyproject.replace(
        NCCL, NCCL_WHL.format(require_nccl_dep) if require_nccl_dep != "na" else ""
    )
    pyproject = (
        f"# Generated by `{os.path.basename(__file__)}`, don't edit: "
        f"'--use-suffix={use_suffix} --require-nccl-dep={require_nccl_dep} "
        f"--create-stub={create_stub}'\n" + pyproject
    )

    with open(OUT_PATH, "w") as fd:
        fd.write(pyproject)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--use-suffix",
        type=str,
        choices=["na", "cpu"] + CUDA_VARIANTS,
        default="na",
        help=(
            "When using this option, rename the package name to xgboost-[suffix]. "
            "Set to 'na' to disable"
        ),
    )
    parser.add_argument(
        "--require-nccl-dep",
        type=str,
        choices=["na"] + CUDA_VARIANTS,
        required=True,
        help="Which NCCL dependency to use; select 'na' to remove NCCL dependency",
    )
    parser.add_argument(
        "--create-stub",
        action="store_true",
        help="Create a stub package that redirects users to install `xgboost`",
    )
    args = parser.parse_args()
    make_pyproject(
        use_suffix=args.use_suffix,
        require_nccl_dep=args.require_nccl_dep,
        create_stub=args.create_stub,
    )
